package datastructure.bplustree;

import java.util.ArrayList;
import java.util.List;

/**
 * B+树实现（阶数为m）
 * 特性：
 * 1. 所有关键字在叶子节点，且叶子节点形成有序链表
 * 2. 非叶子节点仅存索引，不存值
 * 3. 插入/删除后保持平衡
 */
public class BPlusTree {
    private int order; // 阶数（每个节点最多m个子节点）
    private BPlusTreeNode root; // 根节点
    private int minKeyCount; // 每个节点最少关键字数量：ceil(m/2) - 1

    public BPlusTree(int order) {
        if (order < 3) {
            throw new IllegalArgumentException("阶数必须≥3");
        }
        this.order = order;
        this.minKeyCount = (order + 1) / 2 - 1; // 向上取整m/2后减1
        this.root = new BPlusTreeNode(true, order); // 初始根为叶子节点
    }
    /**
     * 查找关键字对应的value
     * @param key 目标关键字
     * @return 对应value，不存在则返回null
     */
    public String search(int key) {
        BPlusTreeNode current = root;
        // 从根节点向下查找，直到叶子节点
        while (!current.isLeaf) {
            int i = 0;
            // 找到第一个大于key的关键字索引
            while (i < current.keys.size() && current.keys.get(i) <= key) {
                i++;
            }
            // 进入对应子节点
            current = current.children.get(i);
        }
        // 在叶子节点中查找key
        int idx = current.keys.indexOf(key);
        return idx != -1 ? current.values.get(idx) : null;
    }
    /**
     * 范围查询：返回[minKey, maxKey]之间的所有value
     */
    public List<String> rangeQuery(int minKey, int maxKey) {
        List<String> result = new ArrayList<>();
        BPlusTreeNode current = root;

        // 定位到第一个可能包含minKey的叶子节点
        while (!current.isLeaf) {
            int i = 0;
            while (i < current.keys.size() && current.keys.get(i) <= minKey) {
                i++;
            }
            current = current.children.get(i);
        }

        // 遍历叶子节点链表，收集符合条件的value
        while (current != null) {
            // 遍历当前叶子节点的keys
            for (int i = 0; i < current.keys.size(); i++) {
                int key = current.keys.get(i);
                if (key > maxKey) {
                    return result; // 超出范围，结束
                }
                if (key >= minKey) {
                    result.add(current.values.get(i));
                }
            }
            current = current.nextLeaf; // 进入下一个叶子节点
        }
        return result;
    }
    /**
     * 插入key-value对
     */
    public void insert(int key, String value) {
        BPlusTreeNode root = this.root;
        // 若根节点已满，先分裂根节点（树高+1）
        if (root.keys.size() == order - 1) {
            BPlusTreeNode newRoot = new BPlusTreeNode(false, order);
            this.root = newRoot;
            newRoot.children.add(root); // 原根节点成为新根的第一个子节点
            splitChild(newRoot, 0); // 分裂新根的第0个子节点
            insertNonFull(newRoot, key, value); // 在新根下插入
        } else {
            insertNonFull(root, key, value); // 根节点未满，直接插入
        }
    }

    /**
     * 向未满的节点插入key-value
     */
    private void insertNonFull(BPlusTreeNode node, int key, String value) {
        int i = node.keys.size() - 1;
        if (node.isLeaf) {
            // 叶子节点：直接插入到有序位置
            // 找到插入位置
            while (i >= 0 && key < node.keys.get(i)) {
                i--;
            }
            node.keys.add(i + 1, key);
            node.values.add(i + 1, value);
        } else {
            // 内部节点：找到子节点并递归插入
            while (i >= 0 && key < node.keys.get(i)) {
                i--;
            }
            i++; // 目标子节点索引
            BPlusTreeNode child = node.children.get(i);
            // 若子节点已满，先分裂
            if (child.keys.size() == order - 1) {
                splitChild(node, i);
                // 分裂后可能需要调整目标子节点
                if (key > node.keys.get(i)) {
                    i++;
                    child = node.children.get(i);
                }
            }
            insertNonFull(child, key, value);
        }
    }

    /**
     * 分裂子节点（将满节点splitNode分裂为两个节点）
     * @param parent 父节点
     * @param childIndex 子节点在父节点中的索引
     */
    private void splitChild(BPlusTreeNode parent, int childIndex) {
        BPlusTreeNode splitNode = parent.children.get(childIndex);
        BPlusTreeNode newNode = new BPlusTreeNode(splitNode.isLeaf, order);
        int mid = minKeyCount; // ceil(order/2) - 1
        int midKey = splitNode.keys.get(mid);

        if (splitNode.isLeaf) {
            // 叶子节点分裂逻辑保持不变
            newNode.keys.addAll(splitNode.keys.subList(mid, splitNode.keys.size()));
            newNode.values.addAll(splitNode.values.subList(mid, splitNode.values.size()));
            splitNode.keys.subList(mid, splitNode.keys.size()).clear();
            splitNode.values.subList(mid, splitNode.values.size()).clear();
            newNode.nextLeaf = splitNode.nextLeaf;
            splitNode.nextLeaf = newNode;
        } else {
            // 内部节点分裂：修复关键字和子节点范围
            newNode.keys.addAll(splitNode.keys.subList(mid + 1, splitNode.keys.size()));
            newNode.children.addAll(splitNode.children.subList(mid + 1, splitNode.children.size()));
            splitNode.keys.subList(mid, splitNode.keys.size()).clear();
            splitNode.children.subList(mid + 1, splitNode.children.size()).clear();
        }

        parent.keys.add(childIndex, midKey);
        parent.children.add(childIndex + 1, newNode);
    }
    /**
     * 删除关键字
     */
    public void delete(int key) {
        deleteHelper(root, key);
        // 若根节点变为空（仅内部节点且无关键字），更新根为其唯一子节点
        if (!root.isLeaf && root.keys.isEmpty()) {
            root = root.children.get(0);
        }
    }
    // 原方法名：findKeyIndex（功能错误）
// 新方法：返回子节点的正确索引（范围 [0, keys.size()]）
    private int findChildIndex(BPlusTreeNode node, int key) {
        int idx = 0;
        // 找到第一个大于等于key的位置，对应子节点索引
        while (idx < node.keys.size() && node.keys.get(idx) < key) {
            idx++;
        }
        return idx; // 即使key不在keys中，也返回正确的子节点索引
    }

    private void deleteHelper(BPlusTreeNode node, int key) {
        // 关键修复：使用新的索引查找方法，获取子节点索引（不会返回-1）
        int childIdx = findChildIndex(node, key);

        if (node.isLeaf) {
            // 叶子节点：仅当索引有效且关键字匹配时删除
            if (childIdx < node.keys.size() && node.keys.get(childIdx) == key) {
                node.keys.remove(childIdx);
                node.values.remove(childIdx);
            }
            return;
        }

        // 内部节点：判断关键字是否在当前节点的keys中
        boolean keyExists = (childIdx < node.keys.size() && node.keys.get(childIdx) == key);
        if (keyExists) {
            // 关键字在当前节点的keys中，处理逻辑不变
            BPlusTreeNode targetChild = node.children.get(childIdx);
            BPlusTreeNode nextChild = node.children.get(childIdx + 1);
            if (targetChild.keys.size() >= minKeyCount + 1) {
                int successorKey = getSuccessor(targetChild);
                deleteHelper(targetChild, successorKey);
                node.keys.set(childIdx, successorKey);
            } else if (nextChild.keys.size() >= minKeyCount + 1) {
                int borrowKey = nextChild.keys.get(0);
                deleteHelper(nextChild, borrowKey);
                node.keys.set(childIdx, borrowKey);
            } else {
                mergeNodes(node, childIdx, targetChild, nextChild);
                deleteHelper(targetChild, key);
            }
        } else {
            // 关键字不在当前节点，去childIdx对应的子节点查找（索引一定有效）
            BPlusTreeNode child = node.children.get(childIdx);
            deleteHelper(child, key);
            // 检查子节点是否下溢
            if (child.keys.size() < minKeyCount) {
                fixUnderflow(node, childIdx);
            }
        }
    }
    // 辅助方法：查找key在节点中的索引
    private int findKeyIndex(BPlusTreeNode node, int key) {
        int idx = 0;
        while (idx < node.keys.size() && node.keys.get(idx) < key) {
            idx++;
        }
        return (idx < node.keys.size() && node.keys.get(idx) == key) ? idx : -1;
    }

    // 获取后继关键字（子节点的最右叶子节点第一个key）
    private int getSuccessor(BPlusTreeNode node) {
        BPlusTreeNode current = node;
        while (!current.isLeaf) {
            current = current.children.get(current.children.size() - 1);
        }
        return current.keys.get(0);
    }

    // 合并两个节点
    private void mergeNodes(BPlusTreeNode parent, int idx, BPlusTreeNode left, BPlusTreeNode right) {
        // 1. 合并关键字（所有节点都需要）
        left.keys.add(parent.keys.get(idx)); // 加入父节点的分隔关键字
        left.keys.addAll(right.keys); // 加入右节点的所有关键字

        // 2. 仅非叶子节点需要合并子节点（叶子节点没有children）
        if (!left.isLeaf) { // 关键修复：判断是否为非叶子节点
            left.children.addAll(right.children);
        }

        // 3. 仅叶子节点需要合并值和维护链表
        if (left.isLeaf) {
            left.values.addAll(right.values);
            left.nextLeaf = right.nextLeaf; // 继承右节点的下一个叶子
        }

        // 4. 父节点移除分隔关键字和右节点
        parent.keys.remove(idx);
        parent.children.remove(idx + 1);
    }

    // 修复节点关键字不足的情况（借调或合并）
    private void fixUnderflow(BPlusTreeNode parent, int childIdx) {
        BPlusTreeNode child = parent.children.get(childIdx);
        // 向左兄弟借调
        if (childIdx > 0) {
            BPlusTreeNode leftSibling = parent.children.get(childIdx - 1);
            if (leftSibling.keys.size() >= minKeyCount + 1) {
                // 借调左兄弟的最后一个key
                child.keys.add(0, parent.keys.get(childIdx - 1));
                if (child.isLeaf) {
                    child.values.add(0, leftSibling.values.remove(leftSibling.values.size() - 1));
                } else {
                    child.children.add(0, leftSibling.children.remove(leftSibling.children.size() - 1));
                }
                parent.keys.set(childIdx - 1, leftSibling.keys.remove(leftSibling.keys.size() - 1));
                return;
            }
        }
        // 向右兄弟借调
        if (childIdx < parent.children.size() - 1) {
            BPlusTreeNode rightSibling = parent.children.get(childIdx + 1);
            if (rightSibling.keys.size() >= minKeyCount + 1) {
                // 借调右兄弟的第一个key
                child.keys.add(parent.keys.get(childIdx));
                if (child.isLeaf) {
                    child.values.add(rightSibling.values.remove(0));
                } else {
                    child.children.add(rightSibling.children.remove(0));
                }
                parent.keys.set(childIdx, rightSibling.keys.remove(0));
                return;
            }
        }
        // 借调失败，合并节点
        if (childIdx > 0) {
            mergeNodes(parent, childIdx - 1, parent.children.get(childIdx - 1), child);
        } else {
            mergeNodes(parent, childIdx, child, parent.children.get(childIdx + 1));
        }
    }
    public static void main(String[] args) {
        BPlusTree bPlusTree = new BPlusTree(3); // 3阶B+树（每个节点最多2个关键字）

        // 插入测试
        bPlusTree.insert(1, "a");
        bPlusTree.insert(3, "b");
        bPlusTree.insert(5, "c");
        bPlusTree.insert(7, "d");
        bPlusTree.insert(9, "e");

        // 查找测试
        System.out.println(bPlusTree.search(3)); // 输出b
        System.out.println(bPlusTree.search(6)); // 输出null

        // 范围查询测试
        System.out.println(bPlusTree.rangeQuery(2, 8)); // 输出[b, c, d]

        // 删除测试
        bPlusTree.delete(3);
        System.out.println(bPlusTree.search(3)); // 输出null
        System.out.println(bPlusTree.rangeQuery(2, 8)); // 输出[c, d]
    }
}